#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jun 29 12:32:23 2021

@author: liangy
E-mail: liangy19@mails.tsinghua.edu.cn

For spherical Jeans equation

Unit System
The unit system we apply here is:
     distance         --> kiloparsec 
     angular distance --> arcsec
     mass             --> solar mass
     time             --> Gyr
     velocity         --> km/s
"""

import numpy as np
from scipy import special
from scipy import constants
from sph_profile import sph_profile as sph

"""
The physical constants required
"""
G   = 4.301e-6        # gravitational constant in (km/s)^2*kpc/M_sun
c   = constants.c*1e-3# speed of light in km/s
arcsec2rad = np.pi/180/3600    # arcsec/rad

def Beta_w(x,a,b):
    a = a + 3e-7  # perturbation to avoid singularities in gamma and betainc
    if a > 0:
        ib = special.betainc(a,b,x)
    else:
        p = int(np.ceil(np.abs(a)))
        tot = 0.0
        for j in range(p):  # don't use gamma recurrence to avoid instabilities
            tot += special.gamma(j+b+a)/special.gamma(j+1+a)*x**(j+a)
        ib = tot*(1-x)**b/special.gamma(b)+special.betainc(a+p,b,x)
    return ib*special.beta(a,b)
    
class sph_jeans():
    """
    The spherical Jeans equation solver by integration with some auxiliary 
    specical function. The integration is accomplished in linear space.
    
    Attention:
        The unit system in this class is kpc-M_sun-km/s, please convert the 
        unit system in advance
        
    Specification (dynamics):
        massless_tracer: bool
            if True, the tracer have no contribution to mass model
            
        polynomial: str, laguerre or legendre
            the option of polynomial applied in integration
            
        npoly: int
            the order of the polynomial
            
        anisotropy: str, const or om
            the option of the anisotropy beta(r) functional form
            
        beta: float, [-inf, 1.0]
            the value of the orbit anisotropy, 
            if anisotropy is om, this quantity will not be used
            
        ra: float, in kpc
            the scale radius of Osipkov-Merritt model
            if anisotropy is const, this quantity will not be used
            
        tracer_model_list: list of str
            the list of the tracer models, usually one or two sersic
            
        tracer_model_para: list of lists of float
            the list including model parameters of tracer, 
            each element is also a list corresponding one of the tracer model
            
        mass_model_list: list of str
            the list of the mass models, usually power law, tracer model + nfw
            
        mass_model_para: list of lists of float
            the list including model parameters of mass/potential
            each element is also a list corresponding one of the mass model
            
        tracer_model_label: list of bool
            the list to decide which model will be accelerated in mass model list
            for example,
            tracer --> [sersic,sersic]
            mass --> [sersic,sersic,two_power_law]
            tracer label --> [True,True,False]
            
            
    """
    def __init__(self,R_,dyn_specification):
        self.R_    = R_   # the radius where we calculate the LOS dispersion
        
        self.dyn_specification = dyn_specification
                
        self.massless_tracer = dyn_specification["massless_tracer"]    # if False, the 
        self.polynomial  = dyn_specification["polynomial"]    # the selection of polynomial used in integration
        self.npoly       = dyn_specification["npoly"]         # the order of the polynomial used in integration
        self.anisotropy  = dyn_specification["anisotropy"]     # anisotropy form, const or om
        
        self.tracer_model_list = dyn_specification["tracer_model_list"]
        self.tracer_model_para = dyn_specification["tracer_model_para"]
        self.mass_model_list   = dyn_specification["mass_model_list"]
        self.mass_model_para   = dyn_specification["mass_model_para"]
        self.tracer_model_label = dyn_specification["tracer_model_label"]
        
    def solver(self):
        """
        spherical Jeans equation solver 
        for calculating the velocity dispersion at the given radius
        """
        if self.polynomial == "laguerre":
            roots, weights = special.roots_laguerre(self.npoly)
            rt = self.R_+roots
            ct = np.exp(roots)
        if self.polynomial == "legendre":
            roots, weights = special.roots_legendre(self.npoly)
            rt = self.R_+(1.0+roots)/(1.0-roots)
            ct = 2.0/(1-roots)**2
        VrmsSigma = np.sum(weights*self.integral_solution(rt,self.R_)*ct)
        Sigma = self.tracer_surface(self.R_)
        Vrms = np.sqrt(VrmsSigma/Sigma)
        return Vrms,Sigma
    
    def accelerated_solver(self,accelerated_nu_roots,accelerated_Mnu_roots):
        """
        spherical Jeans equation solver, accelerated during mcmc sampling
        for calculating the velocity dispersion at the given radius
        """
        if self.polynomial == "laguerre":
            roots, weights = special.roots_laguerre(self.npoly)
            rt = self.R_+roots
            ct = np.exp(roots)
        if self.polynomial == "legendre":
            roots, weights = special.roots_legendre(self.npoly)
            rt = self.R_+(1.0+roots)/(1.0-roots)
            ct = 2.0/(1-roots)**2
            
        if self.massless_tracer:
            ac_integral_solution = self.rest_integral_solution(rt,self.R_)*accelerated_nu_roots*self.enclosed_mass(rt)
            VrmsSigma = np.sum(weights*ac_integral_solution*ct)
        else:
            ac_integral_solution = self.rest_integral_solution(rt,self.R_)*accelerated_nu_roots*(accelerated_Mnu_roots+self.halo_mass(rt))
            VrmsSigma = np.sum(weights*ac_integral_solution*ct)
        Sigma = self.tracer_surface(self.R_)
        Vrms = np.sqrt(VrmsSigma/Sigma)
        return Vrms,Sigma
        
    def integral_solution(self,r,R):
        """
        the integrated function of spherical Jeans equation, r > R
        """
        # choose the anisotropy model
        if self.anisotropy == "const":
            self.beta = self.dyn_specification["beta"]    # constant anisotropy value, [-inf,1]
            return 2*G*self.tracer_density(r)*self.enclosed_mass(r)*self.F_u_const(r/R,self.beta)/r
        if self.anisotropy == "om":
            self.ra = self.dyn_specification["ra"]
            return 2*G*self.tracer_density(r)*self.enclosed_mass(r)*self.F_u_om(r/R,self.ra/R)/r
        
    def rest_integral_solution(self,r,R):
        """
        the model-independent part of integrated function of spherical Jeans equation
        """
        # choose the anisotropy model
        if self.anisotropy == "const":
            self.beta = self.dyn_specification["beta"]    # constant anisotropy value, [-inf,1]
            return 2*G*self.F_u_const(r/R,self.beta)/r
        if self.anisotropy == "om":
            self.ra = self.dyn_specification["ra"]
            return 2*G*self.F_u_om(r/R,self.ra/R)/r
    
    def tracer_density(self,r):
        sph_dyn = sph(r,dim = "3D")
        nu = 0.0
        for i,dyn_nu in enumerate(self.tracer_model_list):
            tracer_para_i = self.tracer_model_para[i]
            if dyn_nu == "power_law":
                nu += sph_dyn.power_law(*tracer_para_i)
            if dyn_nu == "two_power_law":
                nu += sph_dyn.two_power_law(*tracer_para_i)
            if dyn_nu == "sersic":
                nu += sph_dyn.sersic(*tracer_para_i)
            if dyn_nu == "einasto":
                nu += sph_dyn.einasto(*tracer_para_i)
        return nu
    
    def tracer_surface(self,R):
        sph_dyn = sph(R,dim = "2D")
        Sigma = 0.0
        for i,dyn_nu in enumerate(self.tracer_model_list):
            tracer_para_i = self.tracer_model_para[i]
            if dyn_nu == "power_law":
                Sigma += sph_dyn.power_law(*tracer_para_i)
            if dyn_nu == "two_power_law":
                Sigma += sph_dyn.two_power_law(*tracer_para_i)
            if dyn_nu == "sersic":
                Sigma += sph_dyn.sersic(*tracer_para_i)
            if dyn_nu == "einasto":
                Sigma += sph_dyn.einasto(*tracer_para_i)
        return Sigma
    
    def enclosed_mass(self,r):
        sph_dyn = sph(r,dim = "3D")
        mass = 0.0
        for i,dyn_mass in enumerate(self.mass_model_list):
            mass_para_i = self.mass_model_para[i]
            if dyn_mass == "power_law":
                mass += sph_dyn.M_power_law(*mass_para_i)
            if dyn_mass == "two_power_law":
                mass += sph_dyn.M_two_power_law(*mass_para_i)
            if dyn_mass == "sersic":
                mass += sph_dyn.M_sersic(*mass_para_i)
            if dyn_mass == "einasto":
                mass += sph_dyn.M_einasto(*mass_para_i)
        return mass
    
    def tracer_mass(self,r):
        sph_dyn = sph(r,dim = "3D")
        mass = 0.0
        for i,dyn_mass in enumerate(self.tracer_model_list):
            tracer_para_i = self.tracer_model_para[i]
            if dyn_mass == "power_law":
                mass += sph_dyn.M_power_law(*tracer_para_i)
            if dyn_mass == "two_power_law":
                mass += sph_dyn.M_two_power_law(*tracer_para_i)
            if dyn_mass == "sersic":
                mass += sph_dyn.M_sersic(*tracer_para_i)
            if dyn_mass == "einasto":
                mass += sph_dyn.M_einasto(*tracer_para_i)
        return mass
    
    def halo_mass(self,r):
        sph_dyn = sph(r,dim = "3D")
        mass = 0.0
        for i,dyn_mass in enumerate(self.mass_model_list):
            if self.tracer_model_label[i]:
                continue 
            mass_para_i = self.mass_model_para[i]
            if dyn_mass == "power_law":
                mass += sph_dyn.M_power_law(*mass_para_i)
            if dyn_mass == "two_power_law":
                mass += sph_dyn.M_two_power_law(*mass_para_i)
            if dyn_mass == "sersic":
                mass += sph_dyn.M_sersic(*mass_para_i)
            if dyn_mass == "einasto":
                mass += sph_dyn.M_einasto(*mass_para_i)
        return mass
    
    @staticmethod
    def F_u_const(u,beta):
        """
        auxiliary function depending on the anisotropy beta(r)
        assume beta(r) = const
        all of the functions below can be found in Mamon & Locas 2005
        """
        w = 1.0/(u*u)    
        if beta == 0.0:
            return np.sqrt(1-w)
        elif beta == 1.0:
            return (0.5*np.pi*u-np.sqrt(1-w)-u*np.arcsin(1/u))*0.5
        elif beta == 0.5:
            return np.arccosh(u)-np.sqrt(1-w)*0.5
        elif beta == -0.5:
            return np.arccosh(u)*w+np.sqrt(1-w)*0.5
        else:
            return 0.5*u**(2*beta-1)*(
                    (1.5-beta)*np.sqrt(np.pi)*special.gamma(beta-0.5)/special.gamma(beta)
                    +Beta_w(w,beta+0.5,0.5)*beta
                    -Beta_w(w,beta-0.5,0.5))
    
    @staticmethod
    def F_u_om(u,ua):
        """
        auxiliary function depending on the anisotropy beta(r)
        assume beta(r) = r**2/(r**2+ra**2) as Osipkov−Merritt model
        all of the functions below can be found in Mamon & Locas 2005
        """
        w  = 1.0/(u*u)
        # wa = (R*R)/(self.ra*self.ra)
        return 1.0/(ua*ua+1.0)*(
                (ua*ua+0.5)/np.sqrt(ua*ua+1.0)*(u*u+ua*ua)/u*np.arctan(np.sqrt((u*u-1.0)/(ua*ua+1.0)))
                -0.5*np.sqrt(1.0-w))
        
    def acceleration(self):
        """
        the acceleration function for solving jeans equation during mcmc sampling
        attention: the accelerated model must be contained in tracer_model_list
        """
        if self.polynomial == "laguerre":
            roots, weights = special.roots_laguerre(self.npoly)
            rt = self.R_+roots
        if self.polynomial == "legendre":
            roots, weights = special.roots_legendre(self.npoly)
            rt = self.R_+(1.0+roots)/(1.0-roots)
            
        if self.massless_tracer:
            accelerated_nu_roots  = self.tracer_density(rt)
            return accelerated_nu_roots
        else:
            accelerated_nu_roots  = self.tracer_density(rt)
            accelerated_Mnu_roots = self.tracer_mass(rt)
            return accelerated_nu_roots,accelerated_Mnu_roots